import cv2
import os
import numpy as np
from ultralytics import YOLO
import threading
from concurrent.futures import ThreadPoolExecutor

class DocumentDetector:
    def __init__(self, model_path='yolo11n.pt', batch_size=8):
        # 修改默认模型路径
        default_model = os.path.join(os.path.dirname(__file__), 'models', 'yolo11n.pt')
        self.model = YOLO(model_path if os.path.exists(model_path) else default_model)
        self._model_lock = threading.Lock()
        self.batch_size = batch_size
        self.debug_count = 0  # 添加调试计数器
        
    def detect_batch(self, images):
        """批量检测文档"""
        with self._model_lock:
            results = self.model(images, verbose=False, conf=0.2)
            document_boxes = []
            
            for i, result in enumerate(results):
                boxes = result.boxes
                candidates = []  # 存储所有候选框
                best_box = None  # 修改变量名，确保一致性
                
                # 打印检测到的所有类别和置信度
                print(f"\n第 {i} 帧检测结果:")
                for box in boxes:
                    cls = int(box.cls)
                    conf = float(box.conf)
                    print(f"类别: {cls}, 置信度: {conf:.2f}")
                    
                    # 修改检测类别，按优先级排序
                    target_classes = {
                        39: {'priority': 1, 'name': 'book'},     # 书本优先级最高
                        73: {'priority': 1, 'name': 'book'},
                        84: {'priority': 1, 'name': 'book'},
                        28: {'priority': 2, 'name': 'document'}, # 文档次之
                        29: {'priority': 2, 'name': 'document'},
                    }
                    
                    if cls in target_classes:
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        area = (x2 - x1) * (y2 - y1)
                        img_area = images[i].shape[0] * images[i].shape[1]
                        area_ratio = area / img_area
                        
                        # 检测皮肤区域
                        roi = images[i][int(y1):int(y2), int(x1):int(x2)]
                        skin_ratio = self._detect_skin_regions(roi)
                        
                        # 只保留皮肤检测条件
                        if skin_ratio < 0.1:  # 皮肤区域不超过10%
                            candidates.append({
                                'box': np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.int32),
                                'area': area,
                                'priority': target_classes[cls]['priority'],
                                'conf': conf,
                                'cls': cls,
                                'skin_ratio': skin_ratio
                            })
                            print(f"添加候选框: {target_classes[cls]['name']}, "
                                  f"置信度 {conf:.2f}, 皮肤占比 {skin_ratio:.2f}")
                
                # 选择最佳候选框
                if candidates:
                    # 首先按优先级排序，然后按面积和置信度
                    best_candidate = max(candidates, key=lambda x: (x['priority'], x['area'] * x['conf']))
                    document_boxes.append(best_candidate['box'])
                    print(f"选择最佳候选框: 类别 {best_candidate['cls']}, 置信度 {best_candidate['conf']:.2f}")
                else:
                    document_boxes.append(None)
                    # 保存调试信息
                    if self.debug_count < 5:  # 确实是没检测到文档的情况
                        debug_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'debug')
                        os.makedirs(debug_dir, exist_ok=True)
                        
                        # 保存原始图片
                        debug_path = os.path.join(debug_dir, f'no_doc_detected_{self.debug_count}.jpg')
                        cv2.imwrite(debug_path, images[i])
                        
                        # 保存检测结果可视化
                        debug_vis_path = os.path.join(debug_dir, f'no_doc_detected_{self.debug_count}_vis.jpg')
                        result_image = result.plot()
                        cv2.imwrite(debug_vis_path, result_image)
                        
                        print(f"\n未检测到文档，已保存调试图片: {debug_path}")
                        print(f"检测结果可视化: {debug_vis_path}")
                        self.debug_count += 1
                
                # 删除这行，因为已经在上面添加了 document_boxes
                # document_boxes.append(best_box)
            
            return document_boxes

    def _detect_skin_regions(self, frame):
        """检测肤色区域"""
        try:
            if frame is None or frame.size == 0:
                return 0.0
                
            # 转换到YCrCb颜色空间
            ycrcb = cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
            
            # 肤色范围
            lower = np.array([0, 133, 77], dtype=np.uint8)
            upper = np.array([255, 173, 127], dtype=np.uint8)
            
            # 创建肤色掩码
            skin_mask = cv2.inRange(ycrcb, lower, upper)
            
            # 计算肤色区域比例
            skin_ratio = np.sum(skin_mask > 0) / (frame.shape[0] * frame.shape[1])
            
            return skin_ratio
            
        except Exception as e:
            print(f"检测肤色区域错误: {str(e)}")
            return 0.0